#!/usr/bin/python
# Splits a fastq file into multiple.
import sys
import os

# parameters.
fasta_file = sys.argv[1]
sam_file = sys.argv[2]
window = int(sys.argv[3])

def get_sub(txt, hit):
	# min.
	start = hit - window
	if start < 0:
		start = 0

	# max.
	stop = hit + window
	if stop >= len(txt):
		stop = len(txt) - 1

	# return substring.
	return txt[start:stop]

# read in filter.
filters = {}
fin = open(sam_file, "rb")
for line in fin:
	if line[0] == "@": continue
	
	tmp = line.split("\t")
	if tmp[2] not in filters:
		filters[tmp[2]] = []
	filters[tmp[2]].append( (int(tmp[3]), tmp[0] ) )
fin.close()
	
# filter fasta.
first = True
fin = open(fasta_file, "rb")
for line in fin:
	# check if we start new entry.
	if line[0] == ">":
		# tally.
		if first != True:
			# See if we keep this read.
			if head in filters:
				for entry in filters[head]:
					print ">%s^^%s^^%i" % (head, entry[1], entry[0])
					print "%s" % get_sub(seq, entry[0])	
		else:
			first = False
			
		# clear.
		head = line.strip().replace(">","")
		seq = ""
		continue
	
	
	# get only seq.
	seq += line.strip()
fin.close()

# get last.
if head in filters:
	print ">%s" % head
	print "%s" % get_sub(seq, filters[head])

